package com.zshan.clinic.common.util.limit;

import lombok.Data;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Arrays;

@Component
@ConditionalOnProperty(prefix = "rate.limit", name = "enabled", havingValue = "true")
@EnableConfigurationProperties(LuaRateLimiter.RateLimitProperties.class)
public class LuaRateLimiter {

    /**
     *  Redis 脚本执行模板
     */
    private final DefaultRedisScript<Long> script;

    /**
     * Redis 操作模板
     */
    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     *  配置属性，从 yml 中读取
     */
    private final RateLimitProperties properties;

    public LuaRateLimiter(RateLimitProperties properties) {
        this.properties = properties;
        this.script = new DefaultRedisScript<>();
        this.script.setScriptText(
                // KEYS: 1=URI请求计数key, 2=IP总请求计数key, 3=黑名单key
                "local uriKey = KEYS[1] " +
                        "local ipKey = KEYS[2] " +
                        "local blacklistKey = KEYS[3] " +

                        "local expire = tonumber(ARGV[1]) " +        // 计数窗口时间（秒）
                        "local uriLimit = tonumber(ARGV[2]) " +      // URI 限流阈值
                        "local ipLimit = tonumber(ARGV[3]) " +       // IP 限流阈值
                        "local blacklistExpire = tonumber(ARGV[4]) " +  // 黑名单 TTL (秒)

                        // 如果已在黑名单，直接返回 -1
                        "if redis.call('exists', blacklistKey) == 1 then return -1 end " +

                        // URI 限流判断
                        "local uriCount = redis.call('incr', uriKey) " +
                        "if uriCount == 1 then redis.call('expire', uriKey, expire) end " +
                        "if uriCount > uriLimit then " +
                        "    redis.call('set', blacklistKey, 1, 'EX', blacklistExpire) " +
                        "    return -1 " +
                        "end " +

                        // IP 总请求限流判断
                        "local ipCount = redis.call('incr', ipKey) " +
                        "if ipCount == 1 then redis.call('expire', ipKey, expire) end " +
                        "if ipCount > ipLimit then " +
                        "    redis.call('set', blacklistKey, 1, 'EX', blacklistExpire) " +
                        "    return -1 " +
                        "end " +

                        "return 1"
        );
        this.script.setResultType(Long.class);
    }

    /**
     * 限流方法
     *
     * @param ip 请求的 IP 地址
     * @param uri 请求的 URI 路径
     * @return true 表示允许访问，false 表示拒绝（已被拉黑或超限）
     */
    public boolean tryAcquire(String ip, String uri) {
        String uriKey = String.format(RATE_LIMIT_URL_KEY, ip, uri);
        String ipKey = String.format(RATE_LIMIT_IP_KEY, ip);
        String blacklistKey = String.format(RATE_LIMIT_BLACKLIST_KEY, ip);

        Long result = redisTemplate.execute(
                script,
                Arrays.asList(uriKey, ipKey, blacklistKey),
                String.valueOf(properties.getExpire()),             // 统计窗口时间（秒）
                String.valueOf(properties.getUriLimit()),            // 每个 IP 对同一个 URI 最大请求数
                String.valueOf(properties.getIpLimit()),             // 每个 IP 最大请求数
                String.valueOf(properties.getBlacklistExpire())      // 黑名单过期时间（秒）
        );

        return result != null && result == 1L;
    }

    // Redis key 模板：按 URL 限流
    public static final String RATE_LIMIT_URL_KEY = "API:RATE_LIMIT:URL:%s:%s";

    // Redis key 模板：按 IP 限流
    public static final String RATE_LIMIT_IP_KEY = "API:RATE_LIMIT:IP:%s";

    // Redis key 模板：IP 黑名单
    public static final String RATE_LIMIT_BLACKLIST_KEY = "API:RATE_LIMIT:BLACKLIST:%s";

    /**
     * 限流配置类，从 application.yml 配置文件读取
     */
    @ConfigurationProperties(prefix = "rate.limit")
    @Data
    public static class RateLimitProperties {
        /**
         * 是否启用限流功能
         */
        private boolean enabled;

        /**
         * 限流统计窗口时间，单位秒，默认 60 秒
         */
        private int expire = 60;

        /**
         * 每个 IP 对同一个 URI 最大请求次数，默认 30 次
         */
        private int uriLimit = 30;

        /**
         * 每个 IP 所有请求最大次数，默认 100 次
         */
        private int ipLimit = 100;

        /**
         * 黑名单过期时间，单位秒，默认 600 秒
         */
        private int blacklistExpire = 600;
    }
}

